# -*- coding: utf-8 -*-
###################################
# create: 2024-08-27
# author: gql
# desc: 获取一期ddl,生成二期ddl
###################################
import io
import logging
import os
import sys
import time

reload(sys)
sys.setdefaultencoding('utf8')

# tbds hive jdbc
jdbc_url = "jdbc:hive2://10.10.10.10:10000/"
jdbc_usr = "xh"
jdbc_pas = "xh@321"
# 基础路径
base_path = "/data/xh/scripts"
record_path = "/data/xh/scripts/record"
src_ddl_path = "/data/xh/scripts/ddl/src"
dst_ddl_path = "/data/xh/scripts/ddl/dst"
# 建表语句是否自定义location
is_custom_hdfs_location = True
hdfs_db_path = 'hdfs://service/data/ods_back'


class Logger(object):
    @staticmethod
    def set_log_file(app):
        log_path = "logs"
        if not os.path.exists(log_path):
            try:
                os.makedirs(log_path)
            except Exception as e:
                print("其他进程已经创建,{}".format(e))
        return "{}/{}.log".format(log_path, app)

    def __init__(self, app, c_level=logging.INFO, f_level=logging.INFO):
        path = Logger.set_log_file(app)
        self.logger = logging.getLogger(path)
        self.logger.setLevel(logging.INFO)
        fmt = logging.Formatter('[%(asctime)s] [%(levelname)s] %(message)s', '%Y-%m-%d %H:%M:%S')
        ch = logging.StreamHandler()
        ch.setFormatter(fmt)
        ch.setLevel(c_level)
        fh = logging.FileHandler(path, encoding='utf-8')
        fh.setFormatter(fmt)
        fh.setLevel(f_level)
        self.logger.addHandler(ch)
        self.logger.addHandler(fh)
        self.info("log path: {}".format(path))

    def info(self, message):
        self.logger.info(message)

    def error(self, message):
        self.logger.error(message)


log = Logger("src_ddl", c_level=logging.INFO, f_level=logging.INFO)


def date_format(fmt='%Y-%m-%d'):
    times = time.time()
    fmt_time = time.strftime(fmt, time.localtime(times))
    return fmt_time


def exec_shell(script):
    log.info("执行语句: " + script)
    p = os.popen(script)
    out = p.readlines()
    return out, p.close()


def make_dirs(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)


# 读取清单文件
def get_manifest_config(manifest_file):
    config = []
    with io.open(manifest_file, 'r', encoding='utf-8') as fr:
        for line in fr:
            line = line.strip()
            if len(line) == 0:
                continue
            fields = line.split("\t")
            db_table = {"db": "ods", "tab": fields[0]}
            config.append(db_table)
    return config


# 获取orc表ddl并写入文件,获取holo表字段
def gen_src_ddl(batch_no, db_name, src_tab_name):
    script = 'beeline -u "{}" -n {} -p {} --showheader=false --outputformat=csv -e "show create table {}.{};"'.format(
        jdbc_url, jdbc_usr, jdbc_pas, db_name, src_tab_name)
    result_list, _ = exec_shell(script)
    idx = 0
    for line in result_list:
        if '[main]' in line:
            idx += 1
    # src建表语句
    result_list[idx] = result_list[idx].replace("`", "")
    src_ddl_file = "{}/{}/{}_{}_{}.sql".format(src_ddl_path, batch_no, batch_no, db_name, src_tab_name)
    with io.open(src_ddl_file, "w+", encoding="utf-8") as fw:
        for line in result_list:
            line = line.strip()[1:-1] + "\n"
            fw.write(unicode(line))
    fw.close()
    return idx, result_list, src_ddl_file


def get_dst_head(idx, result_list, src_db_name, src_tab_name, dst_db_name, dst_tab_name):
    dst_ddl_head = []
    result_list[idx] = result_list[idx].replace(
        "{}.{}".format(src_db_name, src_tab_name),
        "IF NOT EXISTS {}.{}".format(dst_db_name, dst_tab_name)
    )
    for i in range(len(result_list)):
        line = result_list[i].strip()[1:-1] + "\n"
        dst_ddl_head.append(line)
        # 建表语句到LOCATION则结束
        if line.startswith("LOCATION"):
            if is_custom_hdfs_location:
                dst_ddl_head.append("  '{}/{}'\n".format(hdfs_db_path, dst_tab_name))
                dst_ddl_head.append(";\n")
            else:
                dst_ddl_head[-1] = ";\n"
            break
    return dst_ddl_head


# 生成dst表ddl
def gen_dst_ddl(dst_ddl_head, batch_no, dst_db_name, dst_tab_name):
    dst_ddl_file = "{}/{}/{}_{}_{}.sql".format(dst_ddl_path, batch_no, batch_no, dst_db_name, dst_tab_name)
    with io.open(dst_ddl_file, "w+", encoding="utf-8") as fw:
        for line in dst_ddl_head:
            fw.write(unicode(line))
        fw.flush()
    return dst_ddl_file


# 验证orc表是否存在
def check_tab_exists(db_name, orc_tab_name):
    script = 'beeline -u "{}" -n {} -p {} --showheader=false --outputformat=csv -e "show columns from {}.{};"'.format(
        jdbc_url, jdbc_usr, jdbc_pas, db_name, orc_tab_name)
    result_list = exec_shell(script)
    if len(result_list) == 0:
        return False
    return True


def main(argv):
    if len(argv) < 3:
        print("请输入两个参数: 批次号 批次清单文件\n比如: python {} 001 001.txt".format(argv[0]))
        exit(2)
    # 批次号
    batch_no = argv[1]
    # 清单文件
    manifest_file = argv[2]
    make_dirs("{}/{}".format(src_ddl_path, batch_no))
    make_dirs("{}/{}".format(dst_ddl_path, batch_no))
    make_dirs(record_path)
    # 统计结果输出文件
    out_file_name = "{}/{}_info_{}.csv".format(record_path, batch_no, date_format("%Y%m%d"))
    err_file_name = "{}/{}_not_exists_{}.csv".format(record_path, batch_no, date_format("%Y%m%d"))

    # 获取到批次下所有库和表
    manifest_list = get_manifest_config(manifest_file)

    fw = io.open(out_file_name, "a+", encoding="utf-8")
    fw_err = io.open(err_file_name, "a+", encoding="utf-8")

    # 遍历批次下所有库表
    for db_table in manifest_list:
        src_db_name = db_table["db"]
        src_tab_name = db_table["tab"]
        log.info("正在执行 {}:{}".format(src_db_name, src_tab_name))
        dst_tab_name = src_tab_name
        dst_db_name = src_db_name + "_back"

        # 验证orc表是否存在
        is_exists = check_tab_exists(src_db_name, src_tab_name)
        if not is_exists:
            msg = u"{},{},{}\n".format(batch_no, src_db_name, src_tab_name)
            fw_err.write(msg)
            fw_err.flush()
            log.error("表不存在: {}".format(msg))
            continue
        # 获取src表ddl并写入文件
        try:
            idx, result_list, src_ddl_file = gen_src_ddl(batch_no, src_db_name, src_tab_name)
            dst_ddl_head = get_dst_head(idx, result_list, src_db_name, src_tab_name, dst_db_name, dst_tab_name)
            gen_dst_ddl(dst_ddl_head, batch_no, dst_db_name, dst_tab_name)
        except Exception as e:
            log.error("HANDLE_SRC_DDL error: {}".format(e.message))
            src_final_line = "{},{},{},{},{}".format(
                src_db_name, src_tab_name, "ERROR", "NO_SRC_DDL_FILE", "ERROR")
            dst_final_line = "{},{},{},{},{}".format(
                dst_db_name, dst_tab_name, "ERROR", "NO_DST_DDL_FILE", "ERROR")
            fw.write(unicode(src_final_line + "," + dst_final_line + "\n"))
            continue
        # 获取src信息完成
        src_final_line = "{},{},{}".format(src_db_name, src_tab_name, src_ddl_file)
        fw.write(unicode(src_final_line + "\n"))
        fw.flush()
        log.info("src_final_line {}:{}:{}".format(src_db_name, src_tab_name, src_final_line))

    fw.close()
    fw_err.close()


if __name__ == '__main__':
    main(sys.argv)
